//
// Created by yk120 on 2024/3/3.
//

#include <nanofaiss/IndexLightWeight.h>

#include <cstring>
#include <random>
#include <queue>
#include <omp.h>
#include <mutex>
#include <vector>
#include <unordered_set>

namespace faiss {
IndexLightWeight::IndexLightWeight(size_t d, float sample_portion)
        : d(d), sample_portion(sample_portion) {
    line_size = sizeof(float) * d;
}

void IndexLightWeight::add(size_t nx, const float* x) {
    codes.resize((ntotal + nx) * line_size);
    memcpy(codes.data() + ntotal * line_size, x, nx * line_size);
    ntotal += nx;
}

void IndexLightWeight::train(size_t k) {
    size_t sample_size = static_cast<size_t>(sample_portion * ntotal);
    cluster.resize(sample_size);
    ivf.resize(sample_size);
    std::vector<std::mutex> cluster_latch(sample_size);
    random_sample(
            ntotal, cluster.data(), sample_size); // random sampled id will keep in the cluster

    // pre-reserve
    for (size_t i = 0; i != sample_size; i++) {
        ivf[i].reserve(ntotal * k / sample_size);
    }

    // find the nearest of every point
    // paralle: must use mutex keep every cluster none-competition
    const float* x = reinterpret_cast<const float*>(codes.data());
#pragma omp parallel for
    for (size_t i = 0; i != ntotal; i++) {
        std::priority_queue<Entity> q; // distance, cluster id
        const float* xi = x + i * d;
        for (size_t j = 0; j != sample_size; j++) {
            int sub = cluster[j];
            const float* xj = x + sub * d;
            float distance = fvec_L2(xi, xj, d);
            if (q.size() < k)
                q.push({distance, j});
            else if (q.top().first > distance) {
                q.pop();
                q.push({distance, j});
            }
        }
        while (!q.empty()) {
            auto top = q.top();
            q.pop();

            {
                idx_t c = top.second;
                std::scoped_lock<std::mutex> latch(cluster_latch[c]);
                ivf[c].push_back(i);
            }

        }
    }
}

void IndexLightWeight::search(
        const float* x,
        size_t nt,
        size_t k,
        float* distances,
        idx_t* idx,
        size_t probe) {
    const float* y = reinterpret_cast<const float*>(codes.data());
#pragma omp parallel for
    for (size_t i = 0; i != nt; i++) {
        const float* xi = x + i * d;
        std::priority_queue<Entity> cluster_queue;

        // find the nearest cluster
        for (size_t c = 0; c < cluster.size(); c++) {
            const float* ci = y + cluster[c] * d;
            float distance = fvec_L2(xi, ci, d);
            if (cluster_queue.size() < probe)
                cluster_queue.push({distance, c});
            else if (cluster_queue.top().first > distance) {
                cluster_queue.pop();
                cluster_queue.push({distance, c});
            }
        }

        // deep into the cluster
        std::priority_queue<Entity> q;

        // judge whether point have existed
        std::unordered_set<idx_t> found;

        while (!cluster_queue.empty()) {
            auto [_, c] = cluster_queue.top();
            cluster_queue.pop();
            for (auto member : ivf[c]) {
                const float* yi = y + member * d;
                float distance = fvec_L2(xi, yi, d);
                if (q.size() < k && found.count(member) == 0) {
                    q.emplace(distance, member);
                    found.insert(member);
                }
                else if (q.top().first > distance && found.count(member) == 0) {
                    found.erase(q.top().second);
                    q.pop();
                    q.emplace(distance, member);
                    found.insert(member);
                }
            }
        }

        // copy the result to the array
        float* i_distance = distances + i * k;
        idx_t* i_idx = idx + i * k;
        while (!q.empty()) {
            auto [d, l] = q.top();
            q.pop();
            *i_distance = d;
            *i_idx = l;
            i_distance++;
            i_idx++;
        }
    }
}

void random_sample(size_t n, idx_t *out, size_t sample_size, int seed) {
    if (n <= 0)
        return;

    std::vector<idx_t> permulation(n);
    for (size_t i = 0; i != n; i++)
        permulation[i] = i;

    std::mt19937 mt(static_cast<unsigned>(seed));
    for (size_t i = 0; i != n - 1; i++) {
        int i2 = i + mt() % (n - i);
        std::swap(permulation[i], permulation[i2]);
    }

    memcpy(out, permulation.data(), sample_size * sizeof(idx_t));
}

float fvec_L2(const float* x, const float* y, size_t d) {
    float res = 0.f;
    for (size_t i = 0; i != d; i++)
        res += (x[i] - y[i]) * (x[i] - y[i]);
    return res;
}

}


